import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import sys

#序列长度=7
timesteps = seq_length = 7 # 时间步
data_dim = 5  #输入数据维度=5

# Open,High,Low,Close,Volume
xy = np.loadtxt('./data-02-stock_daily.csv', delimiter=',')
xy = xy[::-1]  # 数据倒序，最后的数据，日期最远的数据

#适当缩放处理
scaler = MinMaxScaler(feature_range=(0, 1))
xy = scaler.fit_transform(xy)

x = xy
y = xy[:, [-1]]  # Close as label y=收盘价

#建立数据集：使用前7天数据=x ,预测第8天的闭盘价格=y
dataX = []
dataY = []
# 使用前7天预测第8天的闭盘价格
for i in range(0, len(y) - seq_length):
    _x = x[i:i + seq_length]
    _y = y[i + seq_length]
    print(_x, "->", _y)
    dataX.append(_x)
    dataY.append(_y)

print('x', tf.shape(dataX))
print('y', tf.shape(dataY))
# sys.exit(0)

#分割数据集：训练集和测试集
# split to train and testing
train_size = int(len(dataY) * 0.7)
test_size = len(dataY) - train_size
trainX, testX = np.array(dataX[0:train_size]), np.array(
    dataX[train_size:len(dataX)])
trainY, testY = np.array(dataY[0:train_size]), np.array(
    dataY[train_size:len(dataY)])

#建立模型
model = Sequential()
model.add(LSTM(1, input_shape=(seq_length, data_dim), return_sequences=False)) # 只有最后返回预测结果
model.add(Dense(1))  # ATTENTION better
# model.add(Activation("linear"))   #线性激活函数  # ATTENTION worse
model.compile(loss='mean_squared_error', optimizer='adam') #mean_squared_error平均绝对值误差

model.summary()

print(trainX.shape, trainY.shape)
model.fit(trainX, trainY, epochs=200)

# make predictions
testPredict = model.predict(testX)

plt.plot(testY)
plt.plot(testPredict)
plt.show()
